import torch
from torch import nn


class ResidualBlock(nn.Module):
    def __init__(self, res_channels, skip_channels, dilation):
        super().__init__()
        self.filter_conv = nn.Conv1d(res_channels, res_channels, 2, dilation=dilation)
        self.gate_conv = nn.Conv1d(res_channels, res_channels, 2, dilation=dilation)
        self.skip_conv = nn.Conv1d(res_channels, skip_channels, 1)
        self.res_conv = nn.Conv1d(skip_channels, res_channels, 1)

    def forward(self, inputs):
        sigmoid_out = torch.sigmoid(self.gate_conv(inputs))  # shape (1, 32, n) -> (1, 32, n - dilation)
        tanh_out = torch.tanh(self.filter_conv(inputs))  # shape (1, 32, n) -> (1, 32, n - dilation)
        mul_out = torch.mul(sigmoid_out, tanh_out)  # shape (1, 32, n - dilation)
        skip_out = self.skip_conv(mul_out)  # shape (1, 32, n - dilation) -> (1, 512, n - dilation)
        res_out = self.res_conv(skip_out)  # shape (1, 512, n - dilation) -> (1, 32, n - dilation)

        # 截取inputs尾部res_out长度的内容, 加上res_out的值
        res_out = res_out + inputs[:, :, -res_out.size(2):]
        return res_out, skip_out  # shape (1, 32, n - dilation), (1, 512, n - dilation)


class WaveNet(nn.Module):
    def __init__(self, in_depth=256, res_channels=32, skip_channels=512, dilation_depth=10, n_repeat=5):
        super().__init__()

        # generate in_depth res_channels——dimensional vectors,
        # and map each incoming number to the corresponding vector.
        # 相当于给传入的数据追加1个长度为res_channels的维度.
        self.preprocess = nn.Embedding(in_depth, res_channels)

        # [1, 2, 4,..., 512, 1, 2, 4,..., 512,..., 512], 共有50个数
        self.dilations = [2 ** i for i in range(dilation_depth)] * n_repeat
        self.main = nn.ModuleList([ResidualBlock(res_channels, skip_channels, dilation) for dilation in self.dilations])
        self.post = nn.Sequential(nn.ReLU(), nn.Conv1d(skip_channels, skip_channels, 1), nn.ReLU(),
                                  nn.Conv1d(skip_channels, in_depth, 1))

    def forward(self, inputs):
        outputs = self.preprocess(inputs).transpose(1, 2)  # shape (1, n) -> (1, n, 32) -> (1, 32, n)
        skip_connections = []
        for layer in self.main:
            outputs, skip = layer(outputs)
            skip_connections.append(skip)
        # 截取最终outputs长度, 加总50层(轮)卷积值. shape (1, 512, m), because skip_channels=512
        outputs = sum([s[:, :, -outputs.size(2):] for s in skip_connections])
        outputs = self.post(outputs)  # shape (1, 512, m) -> (1, 256, m)
        return outputs
